package 力扣_排序算法.快排;

import java.util.Comparator;
import java.util.PriorityQueue;
import java.util.Random;

/**
 * @author zx
 * @create 2022-08-11 16:45
 */
public class 面试题最小K个数_17_14 {
    public int[] smallestK2(int[] arr, int k) {
        if(arr.length == 0 || k == 0){
            return new int[0];
        }
        Random rnd = new Random();
        partition(arr,0,arr.length - 1,rnd);
        int[] res = new int[k];
        for(int i = 0;i < k;i++){
            res[i] = arr[i];
        }
        return res;
    }
    private void partition(int[] arr,int l,int r,Random rnd){
        if(r - l <= 0){
            return;
        }
        int p = l + rnd.nextInt(r - l + 1);
        swap(arr,l,p);
        int v = arr[l];
        int lt = l;
        int i = lt + 1;
        int gt = r + 1;
        while(i < gt){
            if(arr[i] < v){
                swap(arr,i,lt + 1);
                i++;
                lt++;
            }else if(arr[i] > v){
                swap(arr,i,gt - 1);
                gt--;
            }else{
                i++;
            }
        }
        swap(arr,l,lt);
        partition(arr,l,lt - 1,rnd);
        partition(arr,gt,r,rnd);
    }
    private void swap(int[] arr,int i,int j){
        int temp = arr[i];
        arr[i] = arr[j];
        arr[j] = temp;
    }

    public int[] smallestK(int[] arr, int k) {
        //最大堆
        PriorityQueue<Integer> queue = new PriorityQueue<>(new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o2 - o1;
            }
        });
        for(int i = 0;i < arr.length;i++){
            if(queue.size() < k){
                queue.offer(arr[i]);
            }else{
                int peek = queue.peek();
                if(arr[i] > peek){
                    continue;
                }else{
                    queue.poll();
                    queue.offer(arr[i]);
                }
            }
        }
        int[] res = new int[k];
        for(int i = 0;i < k;i++){
            res[i] = queue.poll();
        }
        return res;
    }
}
